import numpy as np
# 直接导入官方ByteTrack
from yolox.tracker.byte_tracker import BYTETracker


class ByteTrackHandler:
    """修复后的跟踪器实现，确保跟踪结果正确返回"""

    def __init__(self, **kwargs):
        # 初始化参数
        self.track_thresh = kwargs.get('track_thresh', 0.5)
        self.track_buffer = kwargs.get('track_buffer', 30)
        self.match_thresh = kwargs.get('match_thresh', 0.8)

        # 初始化官方跟踪器
        self.tracker = BYTETracker(
        )

        # ROI配置
        self.use_roi = kwargs.get('use_roi', False)
        self.roi_rect = kwargs.get('roi_rect', None)
        self.frame_id = 0  # 新增：跟踪帧ID，确保连续

    def set_roi(self, use_roi, roi_rect=None):
        """更新ROI设置"""
        self.use_roi = use_roi
        self.roi_rect = roi_rect
        self.reset()

    def reset(self):
        """重置跟踪器状态（关键修复）"""
        self.tracker = BYTETracker(
        )
        self.frame_id = 0  # 重置帧ID

    def update(self, detections, class_ids):
        """
        修复跟踪更新逻辑，确保返回正确格式的跟踪结果
        输入:
            detections: 检测框数组 [N, 5] (x1,y1,x2,y2,score)
            class_ids: 类别ID数组 [N]
        输出:
            跟踪结果数组 [M, 7] (x1,y1,x2,y2,track_id,class_id,score)
        """
        self.frame_id += 1  # 递增帧ID

        # 1. ROI过滤（修复：确保正确过滤）
        if self.use_roi and self.roi_rect is not None:
            rx1, ry1, rx2, ry2 = self.roi_rect
            valid_indices = []

            for i, det in enumerate(detections):
                x1, y1, x2, y2 = det[:4]
                # 计算检测框与ROI的交并比，确保目标大部分在ROI内
                intersection = max(0, min(x2, rx2) - max(x1, rx1)) * max(0, min(y2, ry2) - max(y1, ry1))
                area = (x2 - x1) * (y2 - y1)
                iou = intersection / (area + 1e-5)

                if iou > 0.5:  # 目标至少50%在ROI内
                    valid_indices.append(i)

            if not valid_indices:
                return np.array([])  # 无有效目标


            detections = detections[valid_indices]
            class_ids = class_ids[valid_indices]


        if len(detections) <= 0:
            #     dets = np.hstack((boxes, scores.reshape(-1, 1)))  # 形状为 (N, 5)
            # else:
            detections = np.empty((0, 5))  # 空检测时传入空数组



        # 3. 调用官方跟踪器更新方法（关键修复：传入正确参数）
        online_targets = self.tracker.update(detections, (1280,720),(1280,720))

        # 4. 整理跟踪结果（修复：确保正确提取跟踪ID和坐标）
        results = []
        for target in online_targets:
            # 不同版本的ByteTrack可能有不同的属性名，这里兼容常见格式
            if hasattr(target, 'tlwh'):
                x1, y1, w, h = target.tlwh
                x2, y2 = x1 + w, y1 + h
            elif hasattr(target, 'xyxy'):
                x1, y1, x2, y2 = target.xyxy
            else:
                continue

            track_id = target.track_id if hasattr(target, 'track_id') else -1
            cls_id = int(target.class_id) if hasattr(target, 'class_id') else 0
            score = target.score if hasattr(target, 'score') else 0.0

            results.append([x1, y1, x2, y2, track_id, cls_id, score])

        return np.array(results)
